package com.netflexitysolutions.amazonws.ec2;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.StringWriter;
import java.io.UnsupportedEncodingException;
import java.net.MalformedURLException;
import java.net.SocketException;
import java.net.URL;
import java.net.URLEncoder;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.text.Collator;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Date;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.TimeZone;
import java.util.TreeMap;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import javax.xml.bind.JAXBException;
import javax.xml.bind.UnmarshalException;

import org.apache.commons.codec.binary.Base64;
import org.apache.commons.httpclient.Header;
import org.apache.commons.httpclient.HostConfiguration;
import org.apache.commons.httpclient.HttpClient;
import org.apache.commons.httpclient.HttpException;
import org.apache.commons.httpclient.HttpMethodBase;
import org.apache.commons.httpclient.MultiThreadedHttpConnectionManager;
import org.apache.commons.httpclient.URI;
import org.apache.commons.httpclient.params.HttpConnectionManagerParams;

public class RestEC2ServiceImpl implements RestEC2Service {

    private String awsAccessId;
    private String awsSecretKey;
    private boolean isSecure;
    private String server;
    private int port;
	private String resourcePrefix = "/";
	private int sigVersion = 2;
	protected Map <String, List<String>> headers;
	// used for caching last used Mac obj.. to save time 99.99% of the time
	private static Map<String, Mac> macMap = new HashMap<String, Mac>();
	private static String lastSecretKey;
	private static Object macSync = new Object();

	private int maxRetries = 5;
	private String userAgent = "typica/";
	private HttpClient hc = null;
	private int maxConnections = 100;
	private String proxyHost = null;
	private int proxyPort;
	private String proxyUser;
	private String proxyPassword;
	private String proxyDomain;	// for ntlm authentication
	private int connectionManagerTimeout = 0;
	private int soTimeout = 0;
	private int connectionTimeout = 0;

    /**
	 * Initializes the queue service with your AWS login information.
	 *
     * @param awsAccessId The your user key into AWS
     * @param awsSecretKey The secret string used to generate signatures for authentication.
     * @param isSecure True if the data should be encrypted on the wire on the way to or from SQS.
     * @param server Which host to connect to.  Usually, this will be s3.amazonaws.com
     * @param port Which port to use.
     */
    public RestEC2ServiceImpl(String awsAccessId, String awsSecretKey, boolean isSecure,
                             String server, int port)
    {
        this.awsAccessId = awsAccessId;
        this.awsSecretKey = awsSecretKey;
        this.isSecure = isSecure;
        this.server = server;
        this.port = port;
		this.headers = new TreeMap<String, List<String>>();
		ArrayList<String> vals = new ArrayList<String>();
		vals.add("2009-04-04");
		this.headers.put("Version", vals);
    }

    /**
     * Make a http request and process the response. This method also performs automatic retries.
	 *
     * @param method The HTTP method to use (GET, POST, DELETE, etc)
     * @param action the name of the action for this query request
     * @param params map of request params
     * @param respType the class that represents the desired/expected return type
     */
	public <T> T makeRequest(HttpMethodBase method, String action, Map<String, String> params, Class<T> respType)
		throws HttpException, IOException, JAXBException, EC2Exception {

		// add auth params, and protocol specific headers
		Map<String, String> qParams = new HashMap<String, String>(params);
		qParams.put("Action", action);
		qParams.put("AWSAccessKeyId", getAwsAccessKeyId());
		qParams.put("SignatureVersion", ""+getSignatureVersion());
		qParams.put("Timestamp", httpDate());
		if (getSignatureVersion() == 2) {
			qParams.put("SignatureMethod", getAlgorithm());
		}
        if (headers != null) {
            for (Iterator<String> i = headers.keySet().iterator(); i.hasNext(); ) {
                String key = i.next();
                for (Iterator<String> j = headers.get(key).iterator(); j.hasNext(); ) {
					qParams.put(key, j.next());
                }
            }
        }
		// sort params by key
		ArrayList<String> keys = new ArrayList<String>(qParams.keySet());
		if (getSignatureVersion() == 2) {
			Collections.sort(keys);
		}
		else {
			Collator stringCollator = Collator.getInstance();
			stringCollator.setStrength(Collator.PRIMARY);
			Collections.sort(keys, stringCollator);
		}

		// build param string
		StringBuilder resource = new StringBuilder();
		if (getSignatureVersion() == 0) {	// ensure Action, Timestamp come first!
			resource.append(qParams.get("Action"));
			resource.append(qParams.get("Timestamp"));
		}
		else if (getSignatureVersion() == 2) {
			resource.append(method.getName());
			resource.append("\n");
			resource.append(getServer().toLowerCase());
			resource.append("\n/");
			String reqURL = makeURL("").toString();
			// see if there is something after the host:port/ in the URL
			if (reqURL.lastIndexOf('/') < (reqURL.length()-1)) {
				// if so, put that here in the string to sign
				// make sure we slice and dice at the right '/'
				int idx = reqURL.lastIndexOf(':');
				resource.append(reqURL.substring(reqURL.indexOf('/', idx)+1));
			}
			resource.append("\n");
			boolean first = true;
			for (String key : keys) {
				if (!first) {
					resource.append("&");
				}
				else { first = false; }
				resource.append(key);
				resource.append("=");
				resource.append(urlencode(qParams.get(key)));
//				System.err.println("encoded params "+key+" :"+(urlencode(qParams.get(key))));
			}
		}
		else {
			for (String key : keys) {
				resource.append(key);
				resource.append(qParams.get(key));
			}
		}
		//System.err.println("String to sign :"+resource.toString());

		// calculate signature
        String encoded = encode(getSecretAccessKey(), resource.toString(), true);

		// build param string, encoding values and adding request signature
		resource = new StringBuilder();
		for (String key : keys) {
			resource.append("&");
			resource.append(key);
			resource.append("=");
			resource.append(urlencode(qParams.get(key)));
		}
		resource.setCharAt(0, '?');	// set first param delimeter
		resource.append("&Signature=");
		resource.append(encoded);

		// finally, build request object
        URL url = makeURL(resource.toString());
		method.setURI(new URI(url.toString(), true));
		method.setRequestHeader(new Header("User-Agent", userAgent));
		if (getSignatureVersion() == 0) {
			method.setRequestHeader(new Header("Content-Type", "application/x-www-form-urlencoded; charset=utf-8"));
		}
		Object response = null;
		boolean done = false;
		int retries = 0;
		boolean doRetry = false;
		EC2Exception error = null;
		do {
			int responseCode = 600;	// default to high value, so we don't think it is valid
			try {
				responseCode = getHttpClient().executeMethod(method);
			} catch (SocketException ex) {
				// these can generally be retried. Treat it like a 500 error
				doRetry = true;
				error = new EC2Exception(ex.getMessage(), ex);
			}
			// 100's are these are handled by httpclient
			if (responseCode < 300) {
				// 200's : parse normal response into requested object
				if (respType != null) {
					InputStream iStr = method.getResponseBodyAsStream();
					response = JAXBuddy.deserializeXMLStream(respType, iStr);
				}
				done = true;
			}
			else if (responseCode < 400) {
				// 300's : what to do?
				throw new HttpException("redirect error : "+responseCode);
			}
			else if (responseCode < 500) {
				// 400's : parse client error message
				String body = getStringFromStream(method.getResponseBodyAsStream());
				throw new EC2Exception(body);
				//throw createException(body, "Client error : ");
			}
			else if (responseCode < 600) {
				// 500's : retry...
				doRetry = true;
				String body = getStringFromStream(method.getResponseBodyAsStream());
				error = new EC2Exception(body);
				//error = createException(body, "");
			}
			if (doRetry) {
				retries++;
				if (retries > maxRetries) {
					throw new HttpException("Number of retries exceeded : "+action, error);
				}
				doRetry = false;
				try { Thread.sleep((int)Math.pow(2.0, retries)*1000); } catch (InterruptedException ex) {}
			}
		} while (!done);
		return (T)response;
	}

	protected HttpClient getHttpClient() {
		if (hc == null) {
			configureHttpClient();
		}
		return hc;
	}

	public void setHttpClient(HttpClient hc) {
		this.hc = hc;
	}

	/**
	 * This method provides the URL for the queue service based on initialization.
	 *
	 * @return generated queue service url
	 */
	public URL getUrl() {
		try {
			return makeURL("");
		} catch (MalformedURLException ex) {
			return null;
		}
	}

	protected String getAwsAccessKeyId() {
		return this.awsAccessId;
	}

	protected String getSecretAccessKey() {
		return this.awsSecretKey;
	}

	protected boolean isSecure() {
		return this.isSecure;
	}

	protected String getServer() {
		return this.server;
	}

	protected void setServer(String server) {
		this.server = server;
	}

	protected int getPort() {
		return this.port;
	}

	public String getResourcePrefix() {
		return this.resourcePrefix;
	}

	public void setResourcePrefix(String prefix) {
		this.resourcePrefix = prefix;
	}

	/**
	 * This method returns the signature version
	 *
	 * @return the version
	 */
	public int getSignatureVersion() {
		return sigVersion;
	}

	/**
	 * This method sets the signature version used to sign requests (0, 1 or 2).
	 * NOTE: This value defaults to 2, so passing 1 is the most likely use case.
	 *
	 * @param version signature version
	 */
	public void setSignatureVersion(int version) {
		if (version != 0 && version != 1 && version != 2) {
			throw new IllegalArgumentException("Only signature versions 0, 1 and 2 supported");
		}
		sigVersion = version;
	}

    /**
     * Create a new URL object for a given resource.
     * @param resource The resource name (bucketName + "/" + key).
     */
    protected URL makeURL(String resource) throws MalformedURLException {
        String protocol = this.isSecure ? "https" : "http";
        return new URL(protocol, this.server, this.port, resourcePrefix+resource);
    }

    /**
     * Calculate the HMAC/SHA1 on a string.
     * @param awsSecretKey passcode to sign it with
     * @param canonicalString data to sign
     * @return signature
     * @throws NoSuchAlgorithmException If the algorithm does not exist.  Unlikely
     * @throws InvalidKeyException If the key is invalid.
     */
    protected String encode(String awsSecretKey, String canonicalString,
                                boolean urlencode) {
		return encode(awsSecretKey, canonicalString, urlencode, getAlgorithm());
	}

    protected String encode(String awsSecretKey, String canonicalString,
                                boolean urlencode, String algorithm) {

        // The following HMAC/SHA1 code for the signature is taken from the
        // AWS Platform's implementation of RFC2104 (amazon.webservices.common.Signature)
        //
        // Acquire an HMAC/SHA1 from the raw key bytes.
        SecretKeySpec signingKey =
            new SecretKeySpec(awsSecretKey.getBytes(), algorithm);

        // Acquire the MAC instance and initialize with the signing key.
		Mac mac = null;
		synchronized (macSync) {
			mac = macMap.get(algorithm);
			if (mac == null || !lastSecretKey.equals(awsSecretKey)) {
				try {
					mac = Mac.getInstance(algorithm);
				} catch (NoSuchAlgorithmException e) {
					// should not happen
					throw new RuntimeException("Could not find sha1 algorithm", e);
				}
				try {
					mac.init(signingKey);
					macMap.put(algorithm, mac);
				} catch (InvalidKeyException e) {
					// also should not happen
					mac = null;
					throw new RuntimeException("Could not initialize the MAC algorithm", e);
				}
				lastSecretKey = awsSecretKey;
			}
		}

        // Compute the HMAC on the digest, and set it.
		byte [] signedBytes = null;
		synchronized (mac) {
			try {
				signedBytes = mac.doFinal(canonicalString.getBytes("UTF-8"));
			} catch (UnsupportedEncodingException e) {
				signedBytes = mac.doFinal(canonicalString.getBytes());
			}
		}
        String b64 = new String(Base64.encodeBase64(signedBytes));

        if (urlencode) {
            return urlencode(b64);
        } else {
            return b64;
        }
    }

	protected String getAlgorithm() {
		return (sigVersion==2)?"HmacSHA256":"HmacSHA1";
	}

    protected String urlencode(String unencoded) {
		String encoded = unencoded;
        try {
			if (sigVersion == 2) {
				encoded = URLEncoder.encode(unencoded, "UTF-8").replace("+", "%20").replace("*", "%2A").replace("%7E", "~");
			}
			else {
            	encoded = URLEncoder.encode(unencoded, "UTF-8");
			}
        } catch (UnsupportedEncodingException e) {
            // should never happen
            throw new RuntimeException("Could not url encode to UTF-8", e);
        }
		return encoded;
    }

	private void configureHttpClient() {
		MultiThreadedHttpConnectionManager connMgr = new MultiThreadedHttpConnectionManager();
		HttpConnectionManagerParams connParams = connMgr.getParams();
		connParams.setMaxTotalConnections(maxConnections);
		connParams.setMaxConnectionsPerHost(HostConfiguration.ANY_HOST_CONFIGURATION, maxConnections);
		connParams.setConnectionTimeout(connectionTimeout);
		connParams.setSoTimeout(soTimeout);
		connMgr.setParams(connParams);
		hc = new HttpClient(connMgr);
// NOTE: These didn't seem to help in my initial testing
//			hc.getParams().setParameter("http.tcp.nodelay", true);
//			hc.getParams().setParameter("http.connection.stalecheck", false); 
		hc.getParams().setConnectionManagerTimeout(connectionManagerTimeout);
		hc.getParams().setSoTimeout(soTimeout);
//		if (proxyHost != null) {
//			HostConfiguration hostConfig = new HostConfiguration();
//			hostConfig.setProxy(proxyHost, proxyPort);
//			hc.setHostConfiguration(hostConfig);
//			log.info("Proxy Host set to "+proxyHost+":"+proxyPort);
//			if (proxyUser != null && !proxyUser.trim().equals("")) {
//				if (proxyDomain != null) {
//					hc.getState().setProxyCredentials(new AuthScope(proxyHost, proxyPort),
//							new NTCredentials(proxyUser, proxyPassword, proxyHost, proxyDomain));
//				}
//				else {
//					hc.getState().setProxyCredentials(new AuthScope(proxyHost, proxyPort),
//							new UsernamePasswordCredentials(proxyUser, proxyPassword));
//				}
//			}
//		}
	}

	private String getStringFromStream(InputStream iStr) throws IOException {
		InputStreamReader rdr = new InputStreamReader(iStr, "UTF-8");
		StringWriter wtr = new StringWriter();
		char [] buf = new char[1024];
		int bytes;
		while ((bytes = rdr.read(buf)) > -1) {
			if (bytes > 0) {
				wtr.write(buf, 0, bytes);
			}
		}
		iStr.close();
		return wtr.toString();
	}

	/**
	 * This method creates a detail packed exception to pass up
	 */
//	private AWSException createException(String errorResponse, String msgPrefix) throws JAXBException {
//		String errorMsg;
//		String requestId;
//		List<AWSError> errors = null;
//		ByteArrayInputStream bais = new ByteArrayInputStream(errorResponse.getBytes());
//		if (errorResponse.indexOf("<ErrorResponse") > -1) {
//			try {
//				// this comes from the SQS2 schema, and is the standard new response
//				ErrorResponse resp = JAXBuddy.deserializeXMLStream(ErrorResponse.class, bais);
//				List<Error> errs = resp.getErrors();
//				errorMsg = "("+errs.get(0).getCode()+") "+errs.get(0).getMessage();
//				requestId = resp.getRequestId();
//				errors = new ArrayList<AWSError>();
//				for (Error e : errs) {
//					errors.add(new AWSError(AWSError.ErrorType.getTypeFromString(e.getType()),
//											e.getCode(), e.getMessage()));
//				}
//			} catch (UnmarshalException ex) {
//				try {
//					// this comes from the DevpayLS schema, duplicated because of the different namespace
//					bais = new ByteArrayInputStream(errorResponse.getBytes());
//					com.xerox.amazonws.typica.jaxb.ErrorResponse resp = JAXBuddy.deserializeXMLStream(com.xerox.amazonws.typica.jaxb.ErrorResponse.class, bais);
//					List<com.xerox.amazonws.typica.jaxb.Error> errs = resp.getErrors();
//					errorMsg = "("+errs.get(0).getCode()+") "+errs.get(0).getMessage();
//					requestId = resp.getRequestID();
//					errors = new ArrayList<AWSError>();
//					for (com.xerox.amazonws.typica.jaxb.Error e : errs) {
//						errors.add(new AWSError(AWSError.ErrorType.getTypeFromString(e.getType()),
//												e.getCode(), e.getMessage()));
//					}
//				} catch (UnmarshalException ex2) {
//					// this comes from the Monitoring schema, duplicated because of the different namespace
//					bais = new ByteArrayInputStream(errorResponse.getBytes());
//					com.xerox.amazonws.typica.monitor.jaxb.ErrorResponse resp = JAXBuddy.deserializeXMLStream(com.xerox.amazonws.typica.monitor.jaxb.ErrorResponse.class, bais);
//					List<com.xerox.amazonws.typica.monitor.jaxb.Error> errs = resp.getErrors();
//					errorMsg = "("+errs.get(0).getCode()+") "+errs.get(0).getMessage();
//					requestId = resp.getRequestId();
//					errors = new ArrayList<AWSError>();
//					for (com.xerox.amazonws.typica.monitor.jaxb.Error e : errs) {
//						errors.add(new AWSError(AWSError.ErrorType.getTypeFromString(e.getType()),
//												e.getCode(), e.getMessage()));
//					}
//				}
//			}
//		}
//		else {
//			Response resp = JAXBuddy.deserializeXMLStream(Response.class, bais);
//			String errorCode = resp.getErrors().getError().getCode();
//			errorMsg = resp.getErrors().getError().getMessage();
//			requestId = resp.getRequestID();
//			if (errorCode != null && !errorCode.trim().equals("")) {
//				errors = new ArrayList<AWSError>();
//				errors.add(new AWSError(AWSError.ErrorType.SENDER, errorCode, errorMsg));
//			}
//		}
//		return new AWSException(msgPrefix + errorMsg, requestId, errors);
//	}

    /**
     * Generate an rfc822 date for use in the Date HTTP header.
     */
    private static String httpDate() {
        final String DateFormat = "yyyy-MM-dd'T'HH:mm:ss'Z'";
        SimpleDateFormat format = new SimpleDateFormat( DateFormat, Locale.US );
        format.setTimeZone( TimeZone.getTimeZone( "GMT" ) );
        return format.format( new Date() );
    }

}
